# -*- coding:utf-8 -*-#
# @Time:2023/7/11 16:05
# @Author:Adong
# @Software:PyCharm


import librosa
import numpy as np
from PIL import Image
from torch import nn
import torch
from sklearn.preprocessing import MinMaxScaler
from torchvision import transforms

class AE(nn.Module):
    '''
    编码器的网络结构
    '''

    def __init__(self, l1_size, l2_size, l3_size, l4_size, l5_size):
        '''

        :param l1_size:第一层尺寸
        :param l2_size:第二层尺寸
        :param l3_size:第三层尺寸
        :param l4_size:第四层尺寸
        :param l5_size:四五层尺寸
        '''
        super(AE, self).__init__()
        self.ReLu = nn.ReLU()  # relu激活函数
        self.linear1 = nn.Linear(l1_size, l2_size)
        self.linear2 = nn.Linear(l2_size, l3_size)
        self.linear3 = nn.Linear(l3_size, l4_size)
        self.linear4 = nn.Linear(l4_size, l5_size)

    def forward(self, x):
        x = self.linear1(x)
        x = self.ReLu(x)
        x = self.linear2(x)
        x = self.linear3(x)
        x = self.ReLu(x)
        x = self.linear4(x)
        x = self.ReLu(x)
        return x

class AEuser:
    def __init__(self,model_weight):
        self.device = "cuda" if torch.cuda.is_available() else 'cpu'
        self.model = AE(l1_size=4096, l2_size=256, l3_size=32, l4_size=256, l5_size=4096).to(self.device)
        self.model.load_state_dict(torch.load(model_weight))  # 加载已训练的网络权重
        self.loss_MSE = torch.nn.MSELoss(reduction='sum')

    def use(self,wav_path,windowlength,stride):
        y,sr = librosa.load(wav_path)
        start = np.linspace(0, len(y) - windowlength, stride)
        piclist = []
        losslist = []
        for s in start:
            piece = y[int(s):int(s + windowlength)]
            pic = piece.reshape(64, 64)
            normalize_tool = MinMaxScaler(feature_range=(0, 255))  # 导入最大最小归一化工具，设置归一化范围为(0,255)
            normpiece = normalize_tool.fit_transform(pic)

            img = Image.fromarray(normpiece).convert('L')
            piclist.append(img)

        for pic in piclist:
            img = transforms.ToTensor()(pic).to(self.device).view(1, 64*64)  # 加载并拉直图片
            re_img = self.model(img)  # 重构
            loss = self.loss_MSE(img, re_img).item()
            losslist.append(loss)
            print("loss MSE = ", loss)

        mean = np.mean(losslist)
        mid = np.median(losslist)
        sigma = np.std(losslist)

        if any(losslist) > mean + 3*sigma or any(losslist) < mean - 2*sigma:
            print("noise")


if __name__ == '__main__':
    weight_path = 'save_model/AutoEncoder_Vdetect.pth'          # 权重路径
    wav_path = 'data/no_noise_wav/直流偏磁_byq_zlpc_01.wav'       # wav路径
    zqd = AEuser(model_weight=weight_path)
    zqd.use(wav_path,4096,2048)